/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package code.google.pigfly;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.metastore.api.FieldSchema;
import org.apache.hadoop.hive.metastore.api.Table;
import org.apache.hadoop.hive.serde2.SerDe;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector.Category;
import org.apache.hadoop.hive.serde2.typeinfo.PrimitiveTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoUtils;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.pig.Expression;
import org.apache.pig.Expression.BinaryExpression;
import org.apache.pig.LoadFunc;
import org.apache.pig.LoadMetadata;
import org.apache.pig.LoadPushDown;
import org.apache.pig.PigException;
import org.apache.pig.ResourceSchema;
import org.apache.pig.ResourceStatistics;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.PigSplit;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.logicalLayer.schema.Schema;
import org.apache.pig.impl.util.UDFContext;

import code.google.pigfly.mapreduce.PigflyInputFormat;
import code.google.pigfly.mapreduce.PigflyRecordReader;

/**
 * A load function that has an ability to read Hive tables into Pig tuples.
 */
public class PigflyLoader extends LoadFunc implements LoadMetadata,
    LoadPushDown {
  public static final Log LOG = LogFactory.getLog(PigflyLoader.class);

  private static final String PARTITION_FILTER = "partition.filter";
  protected static final String PRUNE_PROJECTION_INFO = "prune.projection.info";

  protected String signature;
  protected SerDe serde;
  private String partitionFilterString;

  protected PigflyRecordReader reader;

  /**
   * Keep the path as is
   */
  @Override
  public String relativeToAbsolutePath(String location, Path curDir)
      throws IOException {
    return location;
  }

  @Override
  public InputFormat getInputFormat() throws IOException {
    return new PigflyInputFormat();
  }

  @Override
  public Tuple getNext() throws IOException {
    // read a record from the input steam
    // and then deserialize it with hive serde
    // if (reader.nextKeyValue()) {
    // Writable value = (Writable) reader.getCurrentValue();
    // }

    // covert serde structure into pig Tuple.
    return null;
  }

  @Override
  public void prepareToRead(RecordReader arg0, PigSplit arg1)
      throws IOException {
    // initialize serde

  }

  @Override
  public void setLocation(String location, Job job) throws IOException {
    if (job.getConfiguration().get("mapred.task.id") == null) {
      // only running on the client side
      PigflyInputFormat.setInputLocation(job, location);
    }

  }

  @Override
  public String[] getPartitionKeys(String location, Job job) throws IOException {
    Table table = PigflyMetaStore.getInstance().getTable(location);
    List<FieldSchema> tablePartitionKeys = table.getPartitionKeys();
    String[] partitionKeys = new String[tablePartitionKeys.size()];
    for (int i = 0; i < tablePartitionKeys.size(); i++) {
      partitionKeys[i] = tablePartitionKeys.get(i).getName();
    }
    return partitionKeys;
  }

  @Override
  public ResourceSchema getSchema(String location, Job job) throws IOException {
    Table table = PigflyMetaStore.getInstance().getTable(location);
    List<org.apache.hadoop.hive.metastore.api.FieldSchema> cols =
        table.getSd().getCols();

    List<org.apache.pig.impl.logicalLayer.schema.Schema.FieldSchema> fieldSchemaList =
        new ArrayList<org.apache.pig.impl.logicalLayer.schema.Schema.FieldSchema>(
            cols.size());

    for (org.apache.hadoop.hive.metastore.api.FieldSchema hiveCol : cols) {
      fieldSchemaList
          .add(new org.apache.pig.impl.logicalLayer.schema.Schema.FieldSchema(
              hiveCol.getName(), getPigType(hiveCol.getType())));
    }

    ResourceSchema schema = new ResourceSchema(new Schema(fieldSchemaList));
    return schema;
  }

  @Override
  public ResourceStatistics getStatistics(String arg0, Job arg1)
      throws IOException {
    return null;
  }

  private String getPartitionExpression(Expression expr) {
    if (expr instanceof BinaryExpression) {
      // call getHCatComparisonString on lhs and rhs, and and join the
      // results with OpType string

      // we can just use OpType.toString() on all Expression types except
      // Equal, NotEqualt since Equal has '==' in toString() and
      // we need '='
      String opStr = null;
      switch (expr.getOpType()) {
      case OP_EQ:
        opStr = " = ";
        break;
      default:
        opStr = expr.getOpType().toString();
      }
      BinaryExpression be = (BinaryExpression) expr;
      return "(" + getPartitionExpression(be.getLhs()) + opStr
          + getPartitionExpression(be.getRhs()) + ")";
    } else {
      // should be a constant or column
      return expr.toString();
    }
  }

  @Override
  public void setPartitionFilter(Expression partitionFilter) throws IOException {
    // convert the partition filter expression into a string expected by
    // hcat and pass it in setLocation()
    partitionFilterString = getPartitionExpression(partitionFilter);

    // store this in the udf context so we can get it later
    storeInUDFContext(signature, PARTITION_FILTER, partitionFilterString);
  }

  @Override
  public List<OperatorSet> getFeatures() {
    return Arrays.asList(LoadPushDown.OperatorSet.PROJECTION);
  }

  @Override
  public RequiredFieldResponse pushProjection(RequiredFieldList fields)
      throws FrontendException {
    LOG.info("RequiredFieldList: " + fields.toString());
    // Store the required fields information in the UDFContext so that we
    // can retrieve it later.
    storeInUDFContext(signature, PRUNE_PROJECTION_INFO, fields);

    // HCat will always prune columns based on what we ask of it - so the
    // response is true
    return new RequiredFieldResponse(true);
  }

  // helper methods
  protected void storeInUDFContext(String signature, String key, Object value) {
    UDFContext udfContext = UDFContext.getUDFContext();
    Properties props =
        udfContext
            .getUDFProperties(this.getClass(), new String[] { signature });
    props.put(key, value);
  }

  @Override
  public void setUDFContextSignature(String signature) {
    this.signature = signature;
  }

  private static byte getPigType(String typeString) throws IOException {
    TypeInfo hiveType = TypeInfoUtils.getTypeInfoFromTypeString(typeString);
    String errMsg;
    if (hiveType.getCategory() == Category.PRIMITIVE) {
      return getPrimitiveHType(hiveType);
    }
    errMsg =
        "Pigfly column type '" + hiveType.toString()
            + "' is not supported in Pig as a column type";
    throw new PigException(errMsg, 1116);
  }

  private static byte getPrimitiveHType(TypeInfo hiveType) {
    switch (((PrimitiveTypeInfo) hiveType).getPrimitiveCategory()) {
    case BOOLEAN:
      return DataType.BOOLEAN;
    case BYTE:
      return DataType.BYTE;
    case DOUBLE:
      return DataType.DOUBLE;
    case FLOAT:
      return DataType.FLOAT;
    case INT:
      return DataType.INTEGER;
    case LONG:
      return DataType.LONG;
    case SHORT:
      return DataType.INTEGER;
    case STRING:
      return DataType.CHARARRAY;
    default:
      throw new TypeNotPresentException(
          ((PrimitiveTypeInfo) hiveType).getTypeName(), null);
    }
  }

}